%% Heat Equation Solution Comparison with Collocation Method
clear all; close all; clc;

%% Parameters
nu = 1;              % Diffusion coefficient
L = 1;               % Domain length
h = 1/20;            % Spatial step size
x = 0:h:L;           % Spatial grid
Nx = length(x);    % Number of interior points

% Time parameters
r_values = 1/(2*h);   % r = nu*dt/h^2 values to test
k_values = r_values*h^2/nu; % Corresponding time steps

% Number of steps to plot
steps_to_plot = [1, 2, 10];

%% Initial condition (triangular function)
phi = @(x) (x>=9/20 & x<1/2).*(20*(x-9/20)) + ...
           (x>=1/2 & x<11/20).*(-20*(x-11/20));
u0 = phi(x(1:end))'; % Interior points only

%% Compute and plot solutions
figure('Position', [100, 100, 1200, 600]);
titles = {'After 1 step', 'After 2 steps', 'After 10 steps'};
methods = {'BTCS r=1/2h', 'Collocation r=1/2h'};

for col = 1:3
    steps = steps_to_plot(col);
    
    % BTCS r=1/(2*h) (first row)
    subplot(2, 3, col);
    r = 1/(2*h);
    k = r*h^2/nu;
    u = u0;
    for n = 1:steps
        u = btcs(u, r, Nx);
    end
    plot(x(1:end), u, 'b-', 'LineWidth', 2);
    title(titles{col});
    ylim([0 0.3]);
    if col == 1
        ylabel(methods{1});
    end
    grid on;
    
    
    % Collocation r=1/(2*h) (second row)
    subplot(2, 3, col+3);
    r = 1/(2*h);
    k = r*h^2/nu;
    u = u0;
    for n = 1:steps
        u = collocation(u, r, Nx);
    end
    plot(x(1:end), u, 'g-', 'LineWidth', 2);
    ylim([0 0.3]);
    if col == 1
        ylabel(methods{2});
    end
    grid on;
    
end

sgtitle('Comparison of BTCS and Collocation Methods for Heat Equation');

%% Function Definitions

function u_new = btcs(u_old, r, Nx)
    % Backward Time Centered Space method
    alpha = r;
    main_diag = (1+2*alpha)*ones(Nx,1);
    off_diag = -alpha*ones(Nx-1,1);
    A = diag(main_diag) + diag(off_diag,1) + diag(off_diag,-1);
    u_new = A\u_old;
end

function u_new = collocation(u_old, r, Nx)
    alpha = r;
    
    h = 1/20; 
    k = alpha*h^2;  % Removed /1 since it's redundant
    
    % Create spatial difference operator matrix
    main_diag = -2*ones(Nx,1);
    off_diag = 1*ones(Nx-1,1);
    Dxx = (1/h^2) * (diag(main_diag) + diag(off_diag,1) + diag(off_diag,-1));
    
    % Butcher tableau for collocation method (1/3 and 1 as nodes)
    c = [1/3; 1];
    A = [5/12, -1/12; 
         3/4,   1/4];
    b = [3/4, 1/4];
    
    I = eye(Nx);
    
    % Set up the system for the stages
    % K1 = Dxx*(u_old + k*(A(1,1)*K1 + A(1,2)*K2))
    % K2 = Dxx*(u_old + k*(A(2,1)*K1 + A(2,2)*K2))
    
    % Combine into one big system:
    M = [I - k*A(1,1)*Dxx, -k*A(1,2)*Dxx;
         -k*A(2,1)*Dxx,     I - k*A(2,2)*Dxx];
    
    rhs = [Dxx*u_old; Dxx*u_old];
    
    % Solve for both stages simultaneously
    K_sol = M \ rhs;
    
    % Extract the stages
    K1 = K_sol(1:Nx);
    K2 = K_sol(Nx+1:end);
    
    % Advance the solution
    u_new = u_old + k*(b(1)*K1) + k*(b(2)*K2);
end